package queue

import (
	"context"
	"sync"
)

// 自己包装 cond，加上 WaitTimeout 方法来控制阻塞超时的并发队列
type ConcurrentBlockQueueV2[T any] struct {
	dequeSignal *MyCond
	enqueSignal *MyCond

	mu      *sync.Mutex
	data    []T
	maxSize int
}

func NewConcurrentBlockQueueV2[T any](maxSize int) *ConcurrentBlockQueueV2[T] {
	// 注意内部的 cond 共用了这把锁
	l := &sync.Mutex{}
	return &ConcurrentBlockQueueV2[T]{
		mu:      l,
		data:    make([]T, 0, maxSize),
		maxSize: maxSize,
		dequeSignal: &MyCond{
			Cond: sync.NewCond(l),
		},
		enqueSignal: &MyCond{
			Cond: sync.NewCond(l),
		},
	}
}

func (c *ConcurrentBlockQueueV2[T]) EnQueue(ctx context.Context, data T) error {
	c.mu.Lock()

	for c.isFull() {
		// 进入阻塞前会释放锁，
		// 如果被成功唤醒，又会自动拿到锁，所以要在 for 循环外面手动解锁，
		// 如果超时，不会自动拿到锁，这里也不需要解锁。
		// 这里被唤醒的 goroutin 仍然要和新的 goroutine 抢锁。
		// 这里和用 channel 实现的原理一样，只是这里把解锁、阻塞唤醒、唤醒后加锁、超时都封装起来了。
		err := c.dequeSignal.WaitTimeout(ctx)
		if err != nil {
			return err
		}
	}

	c.data = append(c.data, data)
	c.mu.Unlock()
	c.enqueSignal.Signal()
	return nil
}
func (c *ConcurrentBlockQueueV2[T]) DeQueue(ctx context.Context) (T, error) {
	c.mu.Lock()

	for c.isEmpty() {
		var t T
		// 逻辑和上面类似。
		err := c.enqueSignal.WaitTimeout(ctx)
		if err != nil {
			return t, err
		}
	}
	data := c.data[0]
	c.data = c.data[1:]
	c.mu.Unlock()
	c.dequeSignal.Signal()
	return data, nil
}

func (c *ConcurrentBlockQueueV2[T]) IsFull() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.isFull()
}
func (c *ConcurrentBlockQueueV2[T]) isFull() bool {
	return len(c.data) == c.maxSize
}
func (c *ConcurrentBlockQueueV2[T]) IsEmpty() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.isEmpty()
}
func (c *ConcurrentBlockQueueV2[T]) isEmpty() bool {
	return len(c.data) == 0
}
func (c *ConcurrentBlockQueueV2[T]) Len() uint64 {
	return uint64(len(c.data))
}

type MyCond struct {
	*sync.Cond
}

// 把阻塞的内容放到 goroutine，用一个 chan 在 goroutine 内外传递阻塞是否被唤醒的信号
// 调用前必须已经拿到锁
func (c *MyCond) WaitTimeout(ctx context.Context) error {
	ch := make(chan struct{})
	go func() {
		c.Wait()
		select {
		case ch <- struct{}{}:
		default:
			// 尽管已经超时，但被唤醒后会自动拿到锁，需要处理解锁
			c.L.Unlock()
			// 已经超时退出了，当前用不到的信号要传给下一个
			c.Signal()
		}
	}()
	select {
	case <-ctx.Done():
		// 没被唤醒前超时，不会自动拿到锁，不需要解锁
		return ctx.Err()
	case <-ch:
		return nil
	}
}
